use std::collections::{BTreeMap, HashMap};
use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;

use crate::cluster::Member;

/// 一致性哈希环
/// 
/// 用于在集群节点之间分布虚拟Actor实例
#[derive(Debug, Clone)]
pub struct ConsistentHashRing {
    /// 哈希环 - 存储虚拟节点哈希值到成员的映射
    ring: BTreeMap<u64, Member>,
    /// 每个物理节点对应的虚拟节点数
    replica_count: usize,
    /// 成员ID到该成员虚拟节点哈希列表的映射
    member_hashes: HashMap<String, Vec<u64>>,
}

impl ConsistentHashRing {
    /// 创建一个新的一致性哈希环
    /// 
    /// # Arguments
    /// 
    /// * `replica_count` - 每个物理节点对应的虚拟节点数量
    pub fn new(replica_count: usize) -> Self {
        Self {
            ring: BTreeMap::new(),
            replica_count,
            member_hashes: HashMap::new(),
        }
    }
    
    /// 为指定成员生成虚拟节点哈希值
    fn generate_node_hashes(&self, member: &Member) -> Vec<u64> {
        let mut hashes = Vec::with_capacity(self.replica_count);
        
        for i in 0..self.replica_count {
            let key = format!("{}#{}", member.id, i);
            let hash = Self::hash(&key);
            hashes.push(hash);
        }
        
        hashes
    }
    
    /// 计算字符串的哈希值
    pub fn hash<T: Hash>(key: &T) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }
    
    /// 添加成员到哈希环
    pub fn add_member(&mut self, member: Member) {
        // 首先查看成员是否已存在
        if self.member_hashes.contains_key(&member.id) {
            self.remove_member(&member.id);
        }
        
        // 生成成员的虚拟节点哈希值
        let hashes = self.generate_node_hashes(&member);
        
        // 将虚拟节点添加到环中
        for &hash in &hashes {
            self.ring.insert(hash, member.clone());
        }
        
        // 记录成员的哈希值
        self.member_hashes.insert(member.id.clone(), hashes);
    }
    
    /// 从哈希环中移除成员
    pub fn remove_member(&mut self, member_id: &str) {
        if let Some(hashes) = self.member_hashes.remove(member_id) {
            for hash in hashes {
                self.ring.remove(&hash);
            }
        }
    }
    
    /// 根据键获取负责的节点
    pub fn get_node(&self, key: &str) -> Option<&Member> {
        if self.ring.is_empty() {
            return None;
        }
        
        let hash = Self::hash(&key);
        
        // 找到大于等于hash的第一个节点
        if let Some((_, member)) = self.ring.range(hash..).next() {
            return Some(member);
        }
        
        // 如果没有找到，就返回第一个节点（环绕）
        self.ring.values().next()
    }
    
    /// 获取哈希环中的所有节点
    pub fn get_all_nodes(&self) -> Vec<&Member> {
        let mut nodes = HashMap::new();
        
        for member in self.ring.values() {
            nodes.insert(&member.id, member);
        }
        
        nodes.into_values().collect()
    }
    
    /// 获取哈希环中的节点数量
    pub fn len(&self) -> usize {
        self.member_hashes.len()
    }
    
    /// 检查哈希环是否为空
    pub fn is_empty(&self) -> bool {
        self.member_hashes.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::SystemTime;
    
    fn create_test_member(id: &str, host: &str, port: u16) -> Member {
        Member {
            id: id.to_string(),
            host: host.to_string(),
            port,
            status: crate::cluster::MemberStatus::Alive,
            last_heartbeat: SystemTime::now(),
            role: Some("test".to_string()),
            metadata: HashMap::new(),
        }
    }
    
    #[test]
    fn test_add_and_get_node() {
        let mut ring = ConsistentHashRing::new(10);
        assert!(ring.is_empty());
        
        let member1 = create_test_member("node1", "localhost", 8000);
        let member2 = create_test_member("node2", "localhost", 8001);
        let member3 = create_test_member("node3", "localhost", 8002);
        
        ring.add_member(member1.clone());
        ring.add_member(member2.clone());
        ring.add_member(member3.clone());
        
        assert_eq!(ring.len(), 3);
        
        // 测试一些键的分配
        for i in 0..100 {
            let key = format!("key{}", i);
            let node = ring.get_node(&key).unwrap();
            assert!(node.id == "node1" || node.id == "node2" || node.id == "node3");
        }
    }
    
    #[test]
    fn test_remove_node() {
        let mut ring = ConsistentHashRing::new(10);
        
        let member1 = create_test_member("node1", "localhost", 8000);
        let member2 = create_test_member("node2", "localhost", 8001);
        
        ring.add_member(member1.clone());
        ring.add_member(member2.clone());
        
        assert_eq!(ring.len(), 2);
        
        ring.remove_member("node1");
        assert_eq!(ring.len(), 1);
        
        // 所有键都应该映射到node2
        for i in 0..10 {
            let key = format!("key{}", i);
            let node = ring.get_node(&key).unwrap();
            assert_eq!(node.id, "node2");
        }
    }
    
    #[test]
    fn test_distribution() {
        let mut ring = ConsistentHashRing::new(100); // 使用更多的虚拟节点以获得更均匀的分布
        
        // 添加5个节点
        for i in 0..5 {
            let member = create_test_member(&format!("node{}", i), "localhost", 8000 + i as u16);
            ring.add_member(member);
        }
        
        // 计算1000个键的分布
        let mut counts = HashMap::new();
        for i in 0..1000 {
            let key = format!("key{}", i);
            let node = ring.get_node(&key).unwrap();
            *counts.entry(node.id.clone()).or_insert(0) += 1;
        }
        
        // 检查分布的均匀性（每个节点应该接近200个键）
        for (_, count) in counts {
            assert!(count > 150 && count < 250, "分布不均匀: {}", count);
        }
    }
} 